#from mindformers.trainer import Trainer
# 初始化预训练任务
#trainer = Trainer(task="text_generation", model="llama_13b", train_dataset="")

#res = trainer.predict(input_data="I love Beijing, because",predict_checkpoint="/home/linweibin/lwb2/model/checkpoint_0.ckpt")

import mindspore
#
mindspore.set_device(device_target="CPU", device_id=0)
## 设置 MindSpore 的执行模式为图模式
mindspore.set_context(mode=mindspore.GRAPH_MODE)
from mindformers import AutoModel, AutoTokenizer, AutoConfig
#
## 设置模型类型名
model_type = "llama2_13b"
## 设置测试的输入文本
input_text = "中国的首都是哪个城市？"
#
## 获取模型默认配置项并按需修改
config = AutoConfig.from_pretrained(model_type)
config.use_past = False
#
#
## 根据配置项和预训练权重实例化模型
ckpt = mindspore.load_checkpoint("/home/linweibin/lwb2/model/checkpoint_0.ckpt")
## 实例化tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_type, local_files_only=True)
#
## 对输入进行tokenizer编码
input_ids = tokenizer(input_text, return_tensors="np")["input_ids"]
## 注意：确保input_ids的形状符合模型的要求，例如增加一个批次维度
input_ids = input_ids[None, :]
## 调用model.generate接口执行增量推理
output = ckpt.generate(input_ids, max_length=128, do_sample=False)
## 解码并打印输出
print(tokenizer.decode(output[0]))  # 假设输出是列表形式，需要取第一个元素解码